/* Copyright (C) 1997-2005, Antony L Hosking.
 * All rights reserved.  */
package RegAlloc;

import FlowGraph.AssemNode;
import Assem.Instr;
import Translate.Temp;
import Translate.Tree.Exp;
import Translate.Tree.Exp.CONST;
import Translate.Tree.Exp.MEM;

import java.util.*;

public class Liveness extends InterferenceGraph {
	
	private class Loop {
        public AssemNode header = null, tail = null;
        public Set<AssemNode> nodes = new LinkedHashSet<AssemNode>();

        public Loop(AssemNode header, AssemNode tail) {
            this.header = header;
            this.tail = tail;
            nodes.add(header);
        }
    }
	
    LinkedHashSet<Move> moves = new LinkedHashSet<Move>();

    public LinkedHashSet<Move> moves() {
        return moves;
    }
    
    private Loop FindNaturalLoop(AssemNode header, AssemNode tail) {
    	Stack<AssemNode> workStack = new Stack<AssemNode>();
    	Loop loop = new Loop(header, tail);
    	AssemNode an;
    	
    	if(header != tail) {
    		loop.nodes.add(tail);
    		workStack.push(tail);
    	}
    	
    	while(!workStack.isEmpty()) {
    		an = workStack.pop();
    		for(AssemNode pred : an.preds) {
    			if(!loop.nodes.contains(pred)) {
    				loop.nodes.add(pred);
    				workStack.push(pred);
    			}
    		}
    	}
    	
    	return loop;
    }
    
    private void ComputeLoopDepth(FlowGraph.AssemFlowGraph flow) {
    	AssemNode safeNode = null;
    	LinkedList<AssemNode> worklist = new LinkedList<AssemNode>();
    	LinkedHashSet<AssemNode> workset = new LinkedHashSet<AssemNode>();
    	HashMap<AssemNode, HashMap<AssemNode, Loop>> loopmap = new HashMap<AssemNode, HashMap<AssemNode, Loop>>();
    	
    	// First Solve the Dominator Flow
    	for (AssemNode an : flow.nodes()) {
    		an.dom.addAll(flow.nodes());
    		
    		// Get the "Entry" nodes into the worklist
    		// These nodes with indegree = 0
    		if(an.inDegree() == 0) {
    			worklist.offer(an);
    			workset.add(an);
    		}
    		
    		safeNode = an;
    	}
    	
    	if(worklist.isEmpty()) {
    		// No "Entry" node? Well, then any node as a start should be OK
    		worklist.offer(safeNode);
    		workset.add(safeNode);
    	}
    	
    	// Worklist is ready. 
		AssemNode node = null;
		while((node = worklist.poll()) != null) {
			workset.remove(node);
			LinkedHashSet<AssemNode> newDom = new LinkedHashSet<AssemNode>();
			LinkedHashSet<AssemNode> tempDom = null;
			newDom.add(node);
			for(AssemNode an : node.preds) {
				if(tempDom == null) {
					tempDom = new LinkedHashSet<AssemNode>(an.dom);
				}
				else {
					tempDom.retainAll(an.dom);
				}
			}
			if(tempDom != null) {
				newDom.addAll(tempDom);
			}
			
			if(!newDom.equals(node.dom)) {
				// In set changed
				// All succs should be add to the work list!
				for(AssemNode an : node.succs) {
					if(!workset.contains(an)){
						worklist.offer(an);
						workset.add(an);
					}
    			}
			}

			node.dom = newDom;
		}
		
		// As long as the flow graph is reducible, we could use DFS retreat depth to estimate
		// the loop depth
		for (AssemNode an : flow.nodes()) {
			if(an.inDegree() == 0)
				continue;
			
			for(AssemNode succ : an.succs) {
				if(an.dom.contains(succ)) {
					if(!loopmap.containsKey(succ)) {
						HashMap<AssemNode, Loop> innermap = new HashMap<AssemNode, Loop>();
						loopmap.put(succ, innermap);
						Loop loop = FindNaturalLoop(succ, an);
						innermap.put(an, loop);
					}
					else {
						HashMap<AssemNode, Loop> innermap = loopmap.get(succ);
						if(!innermap.containsKey(an)) {
							Loop loop = FindNaturalLoop(succ, an);
							innermap.put(an, loop);
						}
					}
				}
			}
		}
		
		// All Natural loops of different header are either disjoint or nested.
		for(HashMap<AssemNode, Loop> innermap : loopmap.values()) {
			for(Loop loop : innermap.values()) {
				for(AssemNode an : loop.nodes) {
					an.loopDepth ++;
				}
			}
		}
    }
    
    private void AddCost(Node n, int loopDepth) {
		if(loopDepth == 0) {
			n.spillCost += 1.00;
		}
		else {
			n.spillCost += 1.00 * Math.pow(10.00, loopDepth);
		}
    }
    
    public Liveness(FlowGraph.AssemFlowGraph flow, Translate.Frame frame) {
        // TODO: Compute the interference graph representing liveness
        // information for the temporaris used and defined by the program represented
        // by flow.
    	
    	// Try order the computation following backward of the CFG, using a worklist
    	LinkedList<AssemNode> worklist = new LinkedList<AssemNode>();
    	LinkedHashSet<AssemNode> workset = new LinkedHashSet<AssemNode>();
    	LinkedHashSet<AssemNode> uninitNodes = new LinkedHashSet<AssemNode>();
    	AssemNode safeNode = null;
    	
    	// Clear the liveIn and liveOut list for each node in flow graph
    	for (AssemNode an : flow.nodes()) {
    		uninitNodes.add(an);
    		an.liveIn.clear();
    		an.liveOut.clear();
    		
    		// Get the "Exit" nodes into the worklist
    		// These nodes with outdegree = 0
    		if(an.outDegree() == 0) {
    			worklist.offer(an);
    			workset.add(an);
    		}
    		
    		safeNode = an;
    	}
    	
    	if(worklist.isEmpty()) {
    		// No "Exit" node? Well, then any node as a start should be OK
    		worklist.offer(safeNode);
    		workset.add(safeNode);
    	}
    	
    	// Worklist is ready. 
		AssemNode node = null;
		while((node = worklist.poll()) != null) {
			workset.remove(node);
			LinkedHashSet<Temp> newLiveOut = new LinkedHashSet<Temp>();
			LinkedHashSet<Temp> newLiveIn = new LinkedHashSet<Temp>();
			for(AssemNode an : node.succs) {
				newLiveOut.addAll(an.liveIn);
			}
			
			newLiveIn.addAll(newLiveOut);
			newLiveIn.removeAll(node.def);
			newLiveIn.addAll(node.use);
			
			if((!newLiveIn.equals(node.liveIn)) || (uninitNodes.contains(node))) {
				// In set changed, or just computed from uninitialized. 
				// All preds should be add to the work list!
				for(AssemNode an : node.preds) {
					if(!workset.contains(an)){
						worklist.offer(an);
						workset.add(an);
					}
    			}
			}

			node.liveIn = newLiveIn;
			node.liveOut = newLiveOut;
			
			if(uninitNodes.contains(node))
				uninitNodes.remove(node);
		}
		
		// Compute Loop Depth for all AssemNodes
		ComputeLoopDepth(flow);
    	
    	// Liveness computation is done, build interference graph
		// First put all machine registers into the graph as precolored node
		for (Temp t : frame.registers()) {
			Node n = this.get(t);
			if(n == null) {
				n = new Node(this, t);
				// Precolored node's color is itself
				n.color = n.temp;
			}
		}
		
		// Set degree for each node
		LinkedHashSet<Temp> precolored = new LinkedHashSet<Temp>(Arrays.asList(frame.registers()));
		
		for (AssemNode an : flow.nodes()) {
			LinkedHashSet<Temp> live = new LinkedHashSet<Temp>(an.liveOut);
			Iterator<Assem.Instr> reverse = an.instrs.descendingIterator();
			while(reverse.hasNext()){
				Assem.Instr instr = reverse.next();
				
				if(instr instanceof Assem.Instr.MOVE) {
					Assem.Instr.MOVE moveInstr = (Assem.Instr.MOVE)instr;
					live.remove(moveInstr.src());
					
					Node src = this.get(moveInstr.src());
					if(src == null)
						src = new Node(this, moveInstr.src());
					
					Node dst = this.get(moveInstr.dst());
					if(dst == null)
						dst = new Node(this, moveInstr.dst());
					
					Move move = new Move(src, dst);
					
					dst.moveList.add(move);
					src.moveList.add(move);
					moves.add(move);
				}
				
				for(Temp t : instr.def)
					live.add(t);
				
				for(Temp d : instr.def) {
					Node dNode = this.get(d);
					if(dNode == null)
						dNode = new Node(this, d);
					
					for(Temp l : live) {
						// Add Edge between d and l
						Node lNode = this.get(l);
						if(lNode == null)
							lNode = new Node(this, l);
						
						if(lNode != dNode) {
							// Only Add once
							if((!lNode.goesTo(dNode)) && (!dNode.goesTo(lNode)))
								this.addEdge(dNode, lNode);
						}
					}
				}
				
				for(Temp t : instr.def) {
					live.remove(t);
					Node n = this.get(t);
					if(n != null) {
						AddCost(n, an.loopDepth);
					}
				}
				
				for(Temp t : instr.use) {
					live.add(t);
					Node n = this.get(t);
					if(n != null) {
						AddCost(n, an.loopDepth);
					}
				}
			}
		}
		
		for(Node n : this.nodes()) {
			// precolored nodes have degree N + K
			if(precolored.contains(n.temp)) {
				n.actualDegree = precolored.size() + this.nodes().size();
			}
			else {
				n.actualDegree = n.degree();
			}
		}
		
		
    }

    public void show(java.io.PrintWriter out) {
        for (Node n : nodes()) {
            out.print(n.temp.toString());
            out.print(": ");
            for (Node s : n.succs) {
                out.print(s.temp.toString());
                out.print(" ");
            }
            out.println();
        }
        for (Move move : moves) {
            out.print(move.dst.temp.toString());
            out.print(" <= ");
            out.println(move.src.temp.toString());
        }
        out.flush();
    }
}
